"""
智慧博物馆文物讲解助手 - 主应用
"""

import streamlit as st

# 导入各个功能模块
from modules.ui_components import load_css, add_bg_from_url
from modules.embeddings import load_embedding_model, load_vector_db
from modules.knowledge_graph import load_knowledge_graph
from modules.ai_client import init_ernie_client
from modules.image_utils import load_image_classifier, get_artifact_sample_images
from modules.pages import show_home_page, show_detail_page
from modules.utils import check_required_files,  init_session_state

def main():
    """主函数"""
    # 初始化页面设置
    st.set_page_config(
        page_title="智慧动物园讲解助手",
        page_icon="🦁",
        layout="wide",
        initial_sidebar_state="expanded"
    )
    
    # 加载CSS样式
    load_css()
    
    # 添加背景图片
    bg_url = "https://img.51miz.com/Element/00/71/98/28/0e2c420f_E719828_1b3da8ce.jpg"
    add_bg_from_url(bg_url)
    
    # 初始化会话状态
    init_session_state()
    
    # 检查必要的文件是否存在
    # if not check_required_files():
    #     st.info("首次运行，正在自动构建知识库...")
    #     rebuild_knowledge_base()
    
    # 加载模型和数据
    # 明文设置API密钥
    if "api_key" not in st.session_state or not st.session_state.api_key:
        # 尝试从环境变量获取API密钥
        import os
        api_key_from_env = os.environ.get("AI_STUDIO_API_KEY", "")
        if api_key_from_env:
            st.session_state.api_key = api_key_from_env
        else:
            # 设置默认API密钥（请替换为您的实际API密钥）
            st.session_state.api_key = "8af5cb2dbafc00bc6e6bdc4b68c6fcee4c6bb50a"
    client = init_ernie_client(st.session_state.api_key)
    embedding_model = load_embedding_model()
    kg_data = load_knowledge_graph()
    index, documents = load_vector_db()
    image_classifier = load_image_classifier()
    artifact_images = get_artifact_sample_images()
    
    # 侧边栏设置
    with st.sidebar:
        st.markdown('<h2 class="sub-title">应用设置与信息</h2>', unsafe_allow_html=True)
        
        with st.expander("ERNIE API设置", expanded=client is None):
            api_key_input = st.text_input(
                "请输入百度AI Studio API密钥", 
                value=st.session_state.api_key,
                type="password",
                help="获取密钥请访问：https://aistudio.baidu.com/account/accessToken"
            )
            
            if api_key_input != st.session_state.api_key:
                st.session_state.api_key = api_key_input
                if api_key_input:
                    with st.spinner("正在验证API密钥..."):
                        new_client = init_ernie_client(api_key_input)
                        if new_client:
                            st.success("✅ API密钥验证成功!")
                            st.rerun()
                        else:
                            st.error("❌ API密钥无效或验证失败")
        
        st.markdown(f"""
        <style>
        .status-container {{
            display: flex;
            flex-wrap: wrap;
            gap: 10px;
            margin-bottom: 20px;
        }}
        .status-item {{
            flex: 1;
            min-width: 80px;
            padding: 10px;
            border-radius: 8px;
            text-align: center;
            font-weight: bold;
        }}
        .status-success {{
            background-color: rgba(48, 140, 108, 0.1);
            color: #308C6C;
        }}
        .status-error {{
            background-color: rgba(140, 61, 43, 0.1);
            color: #8C3D2B;
        }}
        .rebuild-button {{
            width: 100%;
            text-align: center;
            margin-top: 15px;
            margin-bottom: 15px;
        }}
        </style>
        
        <div class="status-container">
            <div class="status-item {kg_data and 'status-success' or 'status-error'}">
                知识图谱 {kg_data and '✓' or '✗'}
            </div>
            <div class="status-item {(index and documents) and 'status-success' or 'status-error'}">
                向量数据库 {(index and documents) and '✓' or '✗'}
            </div>
            <div class="status-item {client and 'status-success' or 'status-error'}">
                ERNIE API {client and '✓' or '✗'}
            </div>
        </div>
        """, unsafe_allow_html=True)
        
        st.markdown('<div class="rebuild-button">', unsafe_allow_html=True)
        # if st.button("重建知识库", key="rebuild_db"):
        #     rebuild_knowledge_base()
        st.markdown('</div>', unsafe_allow_html=True)
        
        st.markdown('<div style="height: 1px; background-color: #E0D1B9; margin: 20px 0;"></div>', unsafe_allow_html=True)
        
        if kg_data:
            st.markdown('<h2 class="sub-title">动物导航</h2>', unsafe_allow_html=True)
            
            artifact_nodes = [node for node in kg_data['nodes'] if node['type'] == 'artifact']
            
            era_mapping = {}
            for node in artifact_nodes:
                era = "未知"
                for edge in kg_data['edges']:
                    if edge['source'] == node['id'] and edge['type'] == 'created_in':
                        for n in kg_data['nodes']:
                            if n['id'] == edge['target']:
                                era = n['name']
                                break
                
                if era not in era_mapping:
                    era_mapping[era] = []
                era_mapping[era].append(node['name'])
            
            for era, artifacts in sorted(era_mapping.items()):
                with st.expander(f"{era}（{len(artifacts)}件）", expanded=(era != "未知")):
                    for artifact in sorted(artifacts):
                        artifact_class = "active" if st.session_state.current_artifact == artifact else ""
                        st.markdown(f'<div class="artifact-nav-item {artifact_class}">{artifact}</div>', unsafe_allow_html=True)
                        if st.button(artifact, key=f"nav_{artifact}", help=f"查看{artifact}详情"):
                            st.session_state.current_artifact = artifact
                            st.session_state.page = "detail"
                            st.rerun()
        
        st.markdown('<div style="height: 1px; background-color: #E0D1B9; margin: 20px 0;"></div>', unsafe_allow_html=True)
        
        st.markdown('<h2 class="sub-title">关于应用</h2>', unsafe_allow_html=True)
        st.markdown("""
        <div style="background-color: white; padding: 15px; border-radius: 12px; box-shadow: 0 2px 5px rgba(0,0,0,0.1);">
        <p>智慧动物园讲解助手基于ERNIE 4.5 Turbo VL模型和PaddlePaddle框架开发，
        结合多模态RAG技术和知识图谱，为您提供中国动物的智能讲解服务。</p>
        
        <p>上传动物图片或选择感兴趣的动物，即可获得专业讲解。</p>
        </div>
        """, unsafe_allow_html=True)
        
        st.markdown('<div style="height: 1px; background-color: #E0D1B9; margin: 20px 0;"></div>', unsafe_allow_html=True)
        
        st.markdown('<p style="text-align: center; color: #666666; font-size: 0.9rem;">© 2025 智慧动物园讲解助手</p>', unsafe_allow_html=True)
    
    # 检查资源是否加载完成
    resources_ready = True
    if client is None:
        resources_ready = False
    if embedding_model is None:
        st.error("⚠️ 词嵌入模型未加载，文本搜索功能将不可用")
        resources_ready = False
    if kg_data is None:
        st.error("⚠️ 知识图谱未加载，相关功能将不可用")
        resources_ready = False
    if index is None or documents is None:
        st.error("⚠️ 向量数据库未加载，文档检索功能将不可用")
        resources_ready = False
    
    if not resources_ready:
        if client is None and embedding_model is not None and kg_data is not None and index is not None and documents is not None:
            st.markdown('<div class="info-box">⚠️ 注意：ERNIE客户端未初始化，文物智能问答功能不可用。其他功能可以正常使用。请在侧边栏设置API密钥。</div>', unsafe_allow_html=True)
        else:
            st.markdown('<div class="highlight-box">⚠️ 请先解决资源加载问题，再使用应用功能</div>', unsafe_allow_html=True)
            return
    
    # 根据当前页面显示内容
    if st.session_state.page == "home":
        show_home_page(
            client, embedding_model, kg_data, index, documents, artifact_images
        )
    else:
        show_detail_page(
            client, embedding_model, 
            kg_data, index, documents, artifact_images
        )

if __name__ == "__main__":
    main()